#!/bin/bash
source ./env.sh##环境变量设置，可以放到env.sh 中，也可以用module load+路径，
lrank=$OMPI_COMM_WORLD_LOCAL_RANK
comm_rank=$OMPI_COMM_WORLD_RANK##MPI获取当前进程的rank（序号）
comm_size=$OMPI_COMM_WORLD_SIZE##MPI获取当前通讯域中通讯器的个数


APP="python3 `pwd`/train.py --batch-size=${3} --a=${2} -j 24 --epochs=1 --dist-url tcp://${1}:34567 --dist-backend nccl --world-size=${comm_size} --rank=${comm_rank} --synthetic /work/DeepLearning/Data/ImageNet-pytorch/"，##dist-url tcp://${1}:34567初始化init_method的方法有两种, 一种是使用TCP进行初始化,tcp后是主节点的IP地址##,##dist-backend nccl ##初始化后端nccl,使用GPU建议用nccl##,##/work/DeepLearning/Data/ImageNet-pytorch/"##需要修改写数据所在目录
#python3 的方式运行 .py文件，后边跟的是参数，batch-size，epochs这些就不说了，--world-size ， --rank这些在 .sh文件内容开头就有提到 。


case ${lrank} in
[0])
  export HIP_VISIBLE_DEVICES=0
  export UCX_NET_DEVICES=mlx5_0:1
  export UCX_IB_PCI_BW=mlx5_0:50Gbs
  echo numactl --cpunodebind=0 --membind=0 ${APP}
  numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[1])
  export HIP_VISIBLE_DEVICES=1
  export UCX_NET_DEVICES=mlx5_1:1
  export UCX_IB_PCI_BW=mlx5_1:50Gbs
  echo numactl --cpunodebind=1 --membind=1 ${APP}
  numactl --cpunodebind=1 --membind=1 ${APP}
  ;;
[2])
  export HIP_VISIBLE_DEVICES=2
  export UCX_NET_DEVICES=mlx5_2:1
  export UCX_IB_PCI_BW=mlx5_2:50Gbs
  echo numactl --cpunodebind=2 --membind=2 ${APP} 
  numactl --cpunodebind=2 --membind=2 ${APP}
  ;;
[3])
  export HIP_VISIBLE_DEVICES=3    #分别指定显卡0，1，2，3运行APP内容
  export UCX_NET_DEVICES=mlx5_3:1
  export UCX_IB_PCI_BW=mlx5_3:50Gbs
  echo numactl --cpunodebind=3 --membind=3 ${APP}
  numactl --cpunodebind=3 --membind=3 ${APP}
  ;;
esac
